#include <stdio.h>

// 定义矩阵维度
#define MATRIX_SIZE 1024
#define BLOCK_SIZE 16  // 16x16的线程块

// 用于检查CUDA错误的辅助函数
#define CHECK_CUDA_ERROR(call) \
    do { \
        cudaError_t error = call; \
        if (error != cudaSuccess) { \
            printf("CUDA 错误在行 %d: %s\n", __LINE__, cudaGetErrorString(error)); \
            exit(1); \
        } \
    } while(0)

// 矩阵加法的核函数
__global__ void matrixAdd(const float *A, const float *B, float *C, int size) {
    // 计算当前线程对应的矩阵元素的行和列
    int row = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;

    // 确保线程不会超出矩阵边界
    if (row < size && col < size) {
        // 计算一维数组中的索引
        int index = row * size + col;
        C[index] = A[index] + B[index];
    }
}

int main() {
    int size = MATRIX_SIZE;
    size_t bytes = size * size * sizeof(float);

    // 分配主机内存
    float *h_A = (float*)malloc(bytes);
    float *h_B = (float*)malloc(bytes);
    float *h_C = (float*)malloc(bytes);
    float *h_Verify = (float*)malloc(bytes);

    // 初始化输入矩阵
    for (int i = 0; i < size * size; i++) {
        h_A[i] = rand() / (float)RAND_MAX;
        h_B[i] = rand() / (float)RAND_MAX;
    }

    // 分配设备内存
    float *d_A, *d_B, *d_C;
    CHECK_CUDA_ERROR(cudaMalloc(&d_A, bytes));
    CHECK_CUDA_ERROR(cudaMalloc(&d_B, bytes));
    CHECK_CUDA_ERROR(cudaMalloc(&d_C, bytes));

    // 将数据从主机复制到设备
    CHECK_CUDA_ERROR(cudaMemcpy(d_A, h_A, bytes, cudaMemcpyHostToDevice));
    CHECK_CUDA_ERROR(cudaMemcpy(d_B, h_B, bytes, cudaMemcpyHostToDevice));

    // 设置网格和块的维度
    dim3 threadsPerBlock(BLOCK_SIZE, BLOCK_SIZE);
    dim3 numBlocks((size + BLOCK_SIZE - 1) / BLOCK_SIZE, 
                   (size + BLOCK_SIZE - 1) / BLOCK_SIZE);

    // 启动核函数
    matrixAdd<<<numBlocks, threadsPerBlock>>>(d_A, d_B, d_C, size);
    
    // 检查核函数执行是否有错误
    CHECK_CUDA_ERROR(cudaGetLastError());
    CHECK_CUDA_ERROR(cudaDeviceSynchronize());

    // 将结果从设备复制回主机
    CHECK_CUDA_ERROR(cudaMemcpy(h_C, d_C, bytes, cudaMemcpyDeviceToHost));

    // 在CPU上验证结果
    printf("验证结果中...\n");
    bool correct = true;
    for (int i = 0; i < size * size; i++) {
        h_Verify[i] = h_A[i] + h_B[i];
        if (fabs(h_C[i] - h_Verify[i]) > 1e-5) {
            printf("验证失败：在索引 %d, GPU: %f, CPU: %f\n", 
                   i, h_C[i], h_Verify[i]);
            correct = false;
            break;
        }
    }

    if (correct) {
        printf("验证成功！所有 %dx%d 个元素都计算正确。\n", size, size);
    }

    // 释放设备内存
    CHECK_CUDA_ERROR(cudaFree(d_A));
    CHECK_CUDA_ERROR(cudaFree(d_B));
    CHECK_CUDA_ERROR(cudaFree(d_C));

    // 释放主机内存
    free(h_A);
    free(h_B);
    free(h_C);
    free(h_Verify);

    return 0;
}
